package liming.ai.deepseek.r1;

import liming.ai.abstracts.MessageContext;
import liming.ai.abstracts.AbstractNet;

import liming.ai.deepseek.DeepSeekMessage;
import liming.tool.file.FileRW;
import liming.tool.json.FileConfig;
import liming.tool.json.ToJson;
import liming.tool.json.ValueOf;
import liming.tool.json.annotation.JsonMark;
import liming.tool.json.annotation.JsonMarkCollection;
import liming.tool.lock.CodeCondition;
import liming.tool.lock.CodeLock;
import liming.tool.net.*;
import liming.tool.runnable.HandRunnable;
import liming.tool.runnable.HaveReturnValueInterface;
import okhttp3.*;
import org.jetbrains.annotations.NotNull;
import org.json.JSONArray;
import org.json.JSONObject;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;
import java.util.List;

public class DeepSeekNet extends AbstractNet {
    private static final MediaType JSON = MediaType.get("application/json");
    private final OkHttpClient client;
    private final DeepSeekConfig config;
    public DeepSeekNet() {
        this.client = NetServer.OKHttpConfig.builder().build();
        this.config = new DeepSeekConfig();
        config.deserialize(FileConfig.getConfig(config.getJsonRootName()));
    }
    public ResponseMessage chat(String query, MessageContext messageContext){
        JSONObject body = buildBody(query, messageContext);
        DeepSeekResponseMessage message = body.optBoolean("stream", false) ? streamChat(buildRequest(body.put("stream", true))) : notStreamChat(buildRequest(body.put("stream", false)));
        if(message.status==200)
            messageContext.addMessage(new DeepSeekMessage(DeepSeekMessage.Role.ASSISTANT,message.body));
        else
            messageContext.removeLastMessage();
        return message;

    }
    public DeepSeekResponseMessage streamChat(Request request) {
        return HandRunnable.run(new HaveReturnValueInterface<DeepSeekResponseMessage>() {
            @Override
            public DeepSeekResponseMessage run() throws Exception {
                CodeCondition condition = new CodeLock().newCondition();
                DeepSeekResponseMessage responseDeepSeekMessage =new DeepSeekResponseMessage(200,  "", null);
                client.newCall(request).enqueue(new Callback() {
                    @Override
                    public void onResponse(@NotNull Call call, @NotNull Response response) throws IOException {
                        if(!response.isSuccessful()){
                            responseDeepSeekMessage.status=response.code();
                            responseDeepSeekMessage.message=response.message();
                            responseDeepSeekMessage.body=null;
                            responseDeepSeekMessage.throwable=null;
                            condition.signal();
                            System.out.println(response);
                            call.cancel();
                            return;
                        }
                        try (BufferedReader reader = new BufferedReader(
                                new InputStreamReader(response.body().byteStream()))) {
                            String line;
                            while ((line = reader.readLine()) != null && responseDeepSeekMessage.status==200) {
                                if(streamMessage(line, responseDeepSeekMessage,condition)){
                                    return;
                                };
                            }
                        }catch (Exception e){
                            System.out.println("触发异常");
                            responseDeepSeekMessage.status=500;
                            responseDeepSeekMessage.message=e.getMessage();
                            responseDeepSeekMessage.body=null;
                            responseDeepSeekMessage.thinking=null;
                            responseDeepSeekMessage.throwable = e;
                            condition.signal();
                        }
                    }

                    @Override
                    public void onFailure(@NotNull Call call, @NotNull IOException e) {
                        responseDeepSeekMessage.status=500;
                        responseDeepSeekMessage.message=e.getMessage();
                        responseDeepSeekMessage.body=null;
                        responseDeepSeekMessage.throwable=e;
                        condition.signal();
                        call.clone();
                    }
                });
                condition.await();
                return responseDeepSeekMessage;
            }

            @Override
            public DeepSeekResponseMessage onException(Exception e) {
                return new DeepSeekResponseMessage(500, e.getMessage(), null,e);
            }
        });
    }
    public DeepSeekResponseMessage notStreamChat(Request request) {
        return HandRunnable.run(new HaveReturnValueInterface<DeepSeekResponseMessage>() {
            @Override
            public DeepSeekResponseMessage run() throws Exception {
                try (Response response = client.newCall(request).execute()) {
                    if (!response.isSuccessful()) {
                        return new DeepSeekResponseMessage(response.code(), response.message(), null);
                    }
                    String string = response.body().string();
                    return new DeepSeekResponseMessage(response.code(), response.message(), string);
                }
            }

            @Override
            public DeepSeekResponseMessage onException(Exception e) {
                return new DeepSeekResponseMessage(500, e.getMessage(), null,e);
            }
        });
    }

    public boolean streamMessage(String message, DeepSeekResponseMessage responseDeepSeekMessage, CodeCondition condition){
        System.out.println("message: \n\t"+message);
        if(message.startsWith("data: ")){
            String substring = message.substring(6);
            if(substring.startsWith("[DONE]")){
                condition.signal();
                return true;
            }
            StreamMessage streamMessage = new StreamMessage();
            streamMessage.deserialize(new JSONObject(substring));
            streamMessage.choices.forEach(choice -> {
                responseDeepSeekMessage.addBody(choice.delta.content);
                responseDeepSeekMessage.addThinking(choice.delta.reasoning_content);
            });
        }
        return false;
    }

    private JSONObject buildBody(String query, MessageContext messageContext){
        messageContext.addMessage(new DeepSeekMessage(DeepSeekMessage.Role.USER,query));
        List<JSONObject> messages = messageContext.getMessages();
        return config.body.serializeBody(config.getDefaultClass()).put("messages",new JSONArray(messages));
    }
    Request buildRequest(JSONObject body){
        return new Request.Builder()
                .url(config.header.apiUrl)
                .method("POST", RequestBody.create(body.toString(), JSON))
                .header("Authorization", "Bearer " + config.header.apiKey)
                .header("Content-Type", "application/json")
                .addHeader("Accept", "application/json")
                .build();
    }

    public static class DeepSeekResponseMessage extends ResponseMessage{
        @JsonMark
        public int status;
        @JsonMark
        public String message;
        @JsonMark
        public String thinking;
        @JsonMark
        public String  body;
        @JsonMark
        public Throwable throwable;

        public DeepSeekResponseMessage() {
            super();
        }

        public DeepSeekResponseMessage(int status, String message, String body){
            this(status, message, body, "");
        }
        public DeepSeekResponseMessage(int status, String message, String body,String thinking){
            this(status, message, body,thinking, null);
        }
        public DeepSeekResponseMessage(int status, String message, String body,Throwable throwable){
            this(status, message, body,null, throwable);
        }
        public DeepSeekResponseMessage(int status, String message, String body,String thinking, Throwable throwable){
            this.status = status;
            this.message = message;
            this.body = body;
            this.thinking = thinking;
            this.throwable = throwable;
        }

        void addBody(String body){
            this.body=ValueOf.valueOf(String.class,this.body) +ValueOf.valueOf(String.class,body);
        }
        void addThinking(String thinking){
            this.thinking=ValueOf.valueOf(String.class,this.thinking) +ValueOf.valueOf(String.class,thinking);
        }
    }

    public static class StreamMessage implements ToJson{
        @JsonMark
        public String id;//该对话的唯一标识符。

        @JsonMark
        @JsonMarkCollection(Choice.class)
        public List<Choice> choices;//模型生成的 completion 的选择列表。

        @JsonMark
        public int created;//创建聊天完成时的 Unix 时间戳（以秒为单位）

        @JsonMark
        public String model;//生成该 completion 的模型名。

        @JsonMark
        public String system_fingerprint;//此指纹表示模型运行时使用的后端配置。

        @JsonMark
        public String object;//对象的类型, 其值为 chat.completion。

        @JsonMark
        public Usage usage;

        public static class Choice implements ToJson{
            @JsonMark
            public String finish_reason;//停止原因。可能的值为 stop、length、tool_calls。
            @JsonMark
            public int index;//该 completion 在模型生成的 completion 的选择列表中的索引。
            @JsonMark
            public Delta delta;//模型生成的 completion 消息。
            @JsonMark
            public String logprobs;//该 choice 的对数概率信息。
            public static class Delta implements ToJson{
                @JsonMark
                public String content;//该 completion 的内容。
                @JsonMark
                public String reasoning_content;//仅适用于 deepseek-reasoner 模型。内容为 assistant 消息中在最终答案之前的推理内容。
                @JsonMark
                public String role;//生成这条消息的角色。
                @JsonMark
                public String tool_calls;
            }
        }
        public static class Usage implements ToJson{
            @JsonMark
            public int completion_tokens;//模型 completion 产生的 token 数。
            @JsonMark
            public int prompt_tokens;//用户 prompt 所包含的 token 数。该值等于 prompt_cache_hit_tokens + prompt_cache_miss_tokens
            @JsonMark
            public int prompt_cache_hit_tokens;//用户 prompt 中，命中上下文缓存的 token 数。
            @JsonMark
            public int prompt_cache_miss_tokens;//用户 prompt 中，未命中上下文缓存的 token 数。
            @JsonMark
            public int total_tokens;//该请求中，所有 token 的数量（prompt + completion）。
            @JsonMark
            public Completion_tokens_details completion_tokens_details;
            @JsonMark
            public Prompt_tokens_details prompt_tokens_details;
            public static class Completion_tokens_details implements ToJson{
                @JsonMark
                int reasoning_tokens;//推理模型所产生的思维链 token 数量
            }
            public static class Prompt_tokens_details implements ToJson{
                @JsonMark
                int cached_tokens;//推理模型所产生的思维链 token 数量
            }
        }
    }
}
